from circuit import *
from module import Module
from optimizer import Optimizer
import random

class AdderTree(Module):

    CLK_PORT_NAME = 'clk'
    RST_PORT_NAME = 'rst_n'
    ADDEND_PORT_NAME = 'a'
    SUM_PORT_NAME = 'psum'

    def __init__(self, name: str, input_wires: dict[str, Wire | dict[int, Wire]], output_wires: dict[str, Wire | dict[int, Wire]], wires: set[Wire], circuits: set[Circuit]) -> None:
        super().__init__(name, input_wires, output_wires, wires, circuits)

        self.__parse_e_and_w()
        self.__delete_regs()
        self.__assign_FAs()

    def optimize(self):
        Optimizer(self).optimize()

    def random_test(self, test_times : int) -> None:
        pass_cnt = 0
        failed_cnt = 0
        MAX_VALUE = 2 ** self.__w - 1
        
        for i in range(test_times):
            input_arr = []
            for j in range(2 ** self.__e):
                input_arr.append(random.randint(0, MAX_VALUE))
            
            self.evaluate_by_num_arr(input_arr)
            
            output = self.output_value()
            if output == sum(input_arr):
                pass_cnt += 1
            else:
                failed_cnt += 1
                # print(f'{input_arr}, {bin(sum(input_arr))}, {bin(output)}')
        
        print(f'{pass_cnt} pass, {failed_cnt} failed')

    def output_value(self) -> int:
        """ 
            Get calculate value by output wires
        """
        s = 0
        sum_bits = self.sum_bits()
        for i in range(len(sum_bits)):
            if sum_bits[i].value == True:
                s += 2 ** i
        return s

    def evaluate_by_num_arr(self, number_arr : list[int]):
        # Check input
        if len(number_arr) != 2 ** self.__e:
            print(f'Evaluate Failed, except {2 ** self.__e} input number, but get {len(number_arr)}')
            return

        inputs : dict[str, bool | dict[int, bool]] = {
            AdderTree.CLK_PORT_NAME : False,
            AdderTree.RST_PORT_NAME : False,
            AdderTree.ADDEND_PORT_NAME : {i : False for i in range((2 ** self.__e) * self.__w)}
        }
        
        # For each input number in input array
        for i in range(len(number_arr)):
            number = number_arr[i]
            # Assign Adder Tree's input wires by each number
            for j in range(self.__w):
                inputs[AdderTree.ADDEND_PORT_NAME][self.__w * i + j] = True if (( (number >> j) & 0x1) == 0x1) else False
               
        self.evaluate(inputs)
    
    def __assign_FAs(self):
        self.__FAs : set[FA] = set()
        for c in self._circuits:
            if isinstance(c, FA):
                self.__FAs.add(c)
    
    def __parse_e_and_w(self):
        """
            exx_wxx_.....
        """
        index = self._name.find('_')
        self.__e = int(self._name[1 : index])

        name = self._name[index + 1:]
        index = name.find('_')
        self.__w = int(name[1 : index])

    def __delete_regs(self):
        # Clear input regs
        for wire in self._input_wires[AdderTree.ADDEND_PORT_NAME].values():
            self.__delete_input_reg(wire)

        # Clear output regs
        for wire in self._output_wires[AdderTree.SUM_PORT_NAME].values():
            self.__delete_output_reg(wire)

    def __delete_input_reg(self, wire : Wire):
        reg : REG = wire.output_circuits[0]
        delete_wire = reg.output_wires['Q']
        
        output_circuits : list[Wire] = delete_wire.output_circuits
        wire.output_circuits = output_circuits

        for c in output_circuits:
            for wire_name in c.input_wires:
                if c.input_wires[wire_name] is delete_wire:
                    c.input_wires[wire_name] = wire

        self._wires.remove(delete_wire)
        self._circuits.remove(reg)

    def __delete_output_reg(self, wire : Wire):
        reg : REG = wire.input_circuit
        delete_wire = reg.input_wires['D']
        
        input_circuit : Circuit = delete_wire.input_circuit
        wire.input_circuit = input_circuit

        for wire_name in input_circuit.output_wires:
            if input_circuit.output_wires[wire_name] is delete_wire:
                input_circuit.output_wires[wire_name] = wire

        self._wires.remove(delete_wire)
        self._circuits.remove(reg)

    def show_optimization_report(self):
        S_opt_cnt = 0
        CO_opt_cnt = 0
        all_opt_cnt = 0
        opt_fa_cnt = 0
        opt_port_cnt = 0

        for fa in self.__FAs:
            if fa.S_opt == True and fa.CO_opt == True:
                all_opt_cnt += 1
                opt_fa_cnt += 1
                opt_port_cnt += 2
            elif fa.S_opt == True and fa.CO_opt == False:
                S_opt_cnt += 1
                opt_fa_cnt += 1
                opt_port_cnt += 1
            elif fa.S_opt == False and fa.CO_opt == True:
                CO_opt_cnt += 1
                opt_fa_cnt += 1
                opt_port_cnt += 1

        print('Optimization Report:')
        print(f'\t- (opt FA/FA): {opt_fa_cnt}/{len(self.__FAs)} = {opt_fa_cnt/len(self.__FAs)}')
        print(f'\t- (opt port/port): {opt_port_cnt}/{2 * len(self.__FAs)} = {opt_port_cnt/(2 * len(self.__FAs))}', end='\n\n')

    def show_message(self) -> None:
        print('Adder Tree Message: ')

        print(f'\t- Module name: {self._name}')

        print(f'\t- Number of addends: {2 ** self.__e}')
        print(f'\t- Width of each addend: {self.__w}')
        print(f'\t- Width of sum: {self.__w + self.__e}')

        print(f'\t- Number of wires: {len(self._wires)}')
        print(f'\t- Number of circuits: {len(self._circuits)}')
        print(f'\t- Number of FA: {len(self.__FAs)}')

        print(f'\t- Number of input wires: {len(self.addend_bits())}')
        print(f'\t- Number of output wires: {len(self.sum_bits())}', end='\n\n')

    def FAs(self) -> set[FA]:
        return self.__FAs

    def addend_bits(self) -> dict[int, Wire]:
        return self._input_wires[AdderTree.ADDEND_PORT_NAME]
    
    def sum_bits(self) -> dict[int, Wire]:
        return self._output_wires[AdderTree.SUM_PORT_NAME]

    def e(self) -> int:
        return self.__e

    def w(self) -> int:
        return self.__w